Multi-Token Prediction Done Right
多 Token 预测的正确优化法：Soft-MTP
🧭 LLM 预训练 · Multi-Token Prediction
MIT
提出 Soft-MTP，通过一行代码改动移除标准 MTP 训练中的噪声残差梯度，只保留 KL 自蒸馏信号，全面提升性能与推理加速
📌 工作参考价值（3 点）
- 若内部训练 LLM 用于生成式推荐（如商品描述、搜索 query 推荐、对话式导购），MTP 是一个可落地的训练技巧。采用 Soft-MTP 的一行改动可同时提升模型困惑度和下游准确率，且无额外计算开销，适合直接加到现有预训练流程。
- 在需要 speculative decoding 加速在线推理的场景（如实时搜索建议、Agent 规划生成），Soft-MTP 训练出的 MTP head 能提供更准确的 draft tokens，显著提升 decoding 接受率和加速比，降低延迟。
- 该 KL gradient-recovery trick 可复用到其他共享输出头的自蒸馏任务，例如知识蒸馏或辅助损失设计，避免直接计算大词表 KL 导致的内存瓶颈，兼容现有 fused cross-entropy kernels，工程实现简单。